import hashlib
import os

from pybloom_live import BloomFilter


class BloomFilterHelper(object):
    max_url_num = 5000  # 最大url存储数量
    # todo 看下容错率设成什么值比较不会在url数量上来后误报率增加
    fault_tolerance_rate = 0.0001  # bloom filter容错率

    def __init__(self):
        self.filename = 'bloomFilter.blm'
        is_exist = os.path.exists(self.filename)  # 判断文件是否存在
        if is_exist:
            self.bf = BloomFilter.fromfile(open(self.filename, 'rb'))  # 存在直接打开 储存在内存中
        else:
            self.bf = BloomFilter(self.max_url_num / self.fault_tolerance_rate,
                                  self.fault_tolerance_rate)  # 新建一个 储存在内存中

    # url判重
    def is_url_new(self, url):
        url_encode_md5 = hashlib.md5(url.encode(encoding='utf-8')).hexdigest()
        if url_encode_md5 in self.bf:
            return False

        else:
            self.bf.add(url_encode_md5)
            return True

    def save_file(self):
        self.bf.tofile(open(self.filename, 'wb'))

    def get_max_url_num(self):
        return self.max_url_num
